#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
SIMION多线程工作空间管理模块

解决多线程运行SIMION时的临时文件冲突问题
为每个线程创建独立的工作目录，避免文件冲突
"""

import os
import shutil
import tempfile
import threading
from pathlib import Path
import platform


class SimionWorkspaceManager:
    """SIMION工作空间管理器"""
    
    def __init__(self, base_iob_path, thread_count=None):
        """
        初始化工作空间管理器
        
        参数:
            base_iob_path: 基础IOB文件路径
            thread_count: 线程数量，默认为CPU核心数
        """
        self.base_iob_path = Path(base_iob_path)
        self.base_dir = self.base_iob_path.parent
        self.iob_name = self.base_iob_path.stem
        self.thread_count = thread_count or os.cpu_count()
        
        # 工作空间目录
        self.workspace_dir = self.base_dir / "simion_workspace"
        self.common_dir = self.workspace_dir / "common"
        self.runs_dir = self.workspace_dir / "runs"
        
        # 线程锁
        self._lock = threading.Lock()
        self._created_dirs = set()
    
    def setup_workspace(self):
        """设置工作空间目录结构"""
        try:
            # 先清理旧的工作空间（如果存在）
            if self.workspace_dir.exists():
                print("清理旧的工作空间...")
                try:
                    shutil.rmtree(self.workspace_dir)
                except PermissionError as e:
                    print(f"无法删除工作空间目录，可能被其他进程占用: {e}")
                    print("尝试强制清理...")
                    self._force_cleanup_workspace()
                except Exception as e:
                    print(f"清理工作空间失败: {e}")
                    print("尝试强制清理...")
                    self._force_cleanup_workspace()
            
            # 创建基础目录
            self.workspace_dir.mkdir(exist_ok=True)
            self.common_dir.mkdir(exist_ok=True)
            self.runs_dir.mkdir(exist_ok=True)
            
            # 复制PA文件到common目录
            self._copy_pa_files()
            
            # 为每个线程创建独立目录
            print(f"正在创建 {self.thread_count} 个线程工作目录...")
            for i in range(1, self.thread_count + 1):
                self._create_thread_directory(i)
            
            # 等待所有文件复制完成，避免异步I/O竞争条件
            self._wait_for_all_files_ready()
            
            print(f"工作空间设置完成，支持 {self.thread_count} 个并行线程")
            return True
            
        except Exception as e:
            print(f"设置工作空间失败: {e}")
            return False
    
    def _wait_for_all_files_ready(self):
        """等待所有线程目录的文件复制完成，避免异步I/O竞争条件"""
        import time
        
        max_wait = 10  # 最多等待10秒
        print("等待所有线程文件复制完成...")
        
        for attempt in range(max_wait * 10):
            all_ready = True
            missing_files = []
            
            for i in range(1, self.thread_count + 1):
                run_dir = self.runs_dir / f"run{i}"
                lua_file = run_dir / f"{self.iob_name}.lua"
                iob_file = run_dir / f"{self.iob_name}.iob"
                
                if not lua_file.exists():
                    all_ready = False
                    missing_files.append(f"run{i}/lua")
                if not iob_file.exists():
                    all_ready = False
                    missing_files.append(f"run{i}/iob")
            
            if all_ready:
                print(f"所有线程文件已准备就绪 (等待了 {attempt * 0.1:.1f} 秒)")
                return True
            
            if attempt % 10 == 0:  # 每秒显示一次进度
                print(f"等待中... 缺少文件: {', '.join(missing_files)}")
            
            time.sleep(0.1)  # 每0.1秒检查一次
        
        # 如果超时，显示警告但继续执行
        print(f"警告: 部分文件未及时复制完成，可能导致部分线程执行失败")
        print(f"缺少的文件: {', '.join(missing_files)}")
        return False
    
    def verify_all_thread_files_ready(self):
        """验证所有线程文件是否已准备就绪，如果未准备好则抛出异常"""
        missing_files = []
        
        for i in range(1, self.thread_count + 1):
            run_dir = self.runs_dir / f"run{i}"
            lua_file = run_dir / f"{self.iob_name}.lua"
            iob_file = run_dir / f"{self.iob_name}.iob"
            
            if not lua_file.exists():
                missing_files.append(f"线程{i}的Lua文件: {lua_file}")
            if not iob_file.exists():
                missing_files.append(f"线程{i}的IOB文件: {iob_file}")
        
        if missing_files:
            error_msg = "以下线程文件未准备好:\n" + "\n".join(missing_files)
            raise FileNotFoundError(error_msg)
        
        print(f"验证完成: 所有 {self.thread_count} 个线程的文件已准备就绪")
        return True
    
    def _force_cleanup_workspace(self):
        """强制清理工作空间"""
        try:
            import time
            import psutil
            
            # 等待一段时间让文件解锁
            time.sleep(2)
            
            # 尝试终止可能占用文件的进程
            for proc in psutil.process_iter(['pid', 'name', 'open_files']):
                try:
                    if proc.info['open_files']:
                        for file_info in proc.info['open_files']:
                            if 'simion_workspace' in file_info.path:
                                print(f"发现占用文件的进程: {proc.info['name']} (PID: {proc.info['pid']})")
                                # 不强制终止，只是记录
                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                    pass
            
            # 再次尝试删除
            if self.workspace_dir.exists():
                shutil.rmtree(self.workspace_dir)
                print("强制清理成功")
            
        except Exception as e:
            print(f"强制清理失败: {e}")
            # 如果还是失败，尝试重命名目录
            try:
                if self.workspace_dir.exists():
                    import datetime
                    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
                    old_workspace = self.workspace_dir.parent / f"simion_workspace_old_{timestamp}"
                    self.workspace_dir.rename(old_workspace)
                    print(f"重命名旧工作空间为: {old_workspace}")
            except Exception as rename_e:
                print(f"重命名工作空间也失败: {rename_e}")
    
    def _copy_pa_files(self):
        """复制PA文件到common目录"""
        pa_files = list(self.base_dir.glob(f"{self.iob_name}.pa*"))
        
        if pa_files:
            print(f"复制 {len(pa_files)} 个PA文件到common目录")
            for pa_file in pa_files:
                target_file = self.common_dir / pa_file.name
                if not target_file.exists():
                    try:
                        shutil.copy2(pa_file, target_file)
                    except Exception as e:
                        print(f"复制PA文件失败 {pa_file.name}: {e}")
    
    def _create_thread_directory(self, thread_id):
        """为指定线程创建独立目录"""
        run_dir = self.runs_dir / f"run{thread_id}"
        run_dir.mkdir(exist_ok=True)
        
        # 创建输出子目录
        output_dir = run_dir / "output"
        output_dir.mkdir(exist_ok=True)
        
        # 复制IOB文件（避免多线程同时打开同一源文件）
        iob_file = run_dir / f"{self.iob_name}.iob"
        if not iob_file.exists():
            try:
                with self._lock:  # 使用线程锁保护文件复制
                    # 先读取源文件内容，再写入目标文件，避免多线程同时打开源文件
                    with open(self.base_iob_path, 'rb') as src_f:
                        data = src_f.read()
                    with open(iob_file, 'wb') as dst_f:
                        dst_f.write(data)
                        os.fsync(dst_f.fileno())
            except Exception as e:
                print(f"复制IOB文件失败 run{thread_id}: {e}")
        
        # 复制FLY2文件（如果存在）
        fly2_file = self.base_dir / f"{self.iob_name}.fly2"
        if fly2_file.exists():
            target_fly2 = run_dir / f"{self.iob_name}.fly2"
            if not target_fly2.exists():
                try:
                    with self._lock:  # 使用线程锁保护文件复制
                        # 先读取源文件内容，再写入目标文件，避免多线程同时打开源文件
                        with open(fly2_file, 'rb') as src_f:
                            data = src_f.read()
                        with open(target_fly2, 'wb') as dst_f:
                            dst_f.write(data)
                            os.fsync(dst_f.fileno())
                except Exception as e:
                    print(f"复制FLY2文件失败 run{thread_id}: {e}")
        
        # 复制Lua文件（如果存在）- 强制覆盖以确保使用最新版本
        lua_file = self.base_dir / f"{self.iob_name}.lua"
        if lua_file.exists():
            target_lua = run_dir / f"{self.iob_name}.lua"
            try:
                with self._lock:  # 使用线程锁保护文件复制
                    # 先读取源文件内容，再写入目标文件，避免多线程同时打开源文件
                    with open(lua_file, 'rb') as src_f:
                        data = src_f.read()
                    with open(target_lua, 'wb') as dst_f:
                        dst_f.write(data)
                        os.fsync(dst_f.fileno())
            except Exception as e:
                print(f"复制Lua文件失败 run{thread_id}: {e}")
        
        # 复制excellib.lua文件（如果存在）
        excellib_file = self.base_dir / "excellib.lua"
        if excellib_file.exists():
            target_excellib = run_dir / "excellib.lua"
            if not target_excellib.exists():
                try:
                    with self._lock:  # 使用线程锁保护文件复制
                        # 先读取源文件内容，再写入目标文件，避免多线程同时打开源文件
                        with open(excellib_file, 'rb') as src_f:
                            data = src_f.read()
                        with open(target_excellib, 'wb') as dst_f:
                            dst_f.write(data)
                            os.fsync(dst_f.fileno())
                except Exception as e:
                    print(f"复制excellib.lua失败 run{thread_id}: {e}")
        
        # 创建PA文件的符号链接（Windows）或硬链接（Linux/Mac）
        self._create_pa_links(run_dir)
    
    def _create_pa_links(self, run_dir):
        """为PA文件创建链接"""
        pa_files = list(self.common_dir.glob(f"{self.iob_name}.pa*"))
        
        for pa_file in pa_files:
            target_link = run_dir / pa_file.name
            if not target_link.exists():
                try:
                    if platform.system() == "Windows":
                        # Windows使用硬链接
                        os.link(str(pa_file), str(target_link))
                    else:
                        # Linux/Mac使用符号链接
                        os.symlink(str(pa_file), str(target_link))
                except Exception as e:
                    print(f"创建PA链接失败 {pa_file.name}: {e}")
    
    def get_thread_workspace(self, thread_id):
        """获取指定线程的工作空间路径"""
        run_dir = self.runs_dir / f"run{thread_id}"
        return {
            'iob_path': str(run_dir / f"{self.iob_name}.iob"),
            'output_dir': str(run_dir / "output"),
            'work_dir': str(run_dir)
        }
    
    def cleanup_workspace(self, keep_workspace=False):
        """
        清理SIMION工作空间
        
        参数:
            keep_workspace (bool): 是否保留目录结构和文件
        """
        try:
            if not self.workspace_dir.exists():
                return
            
            if keep_workspace:
                print("保留工作空间结构，仅清理临时文件...")
                # 保留所有 run 目录、pa、iob、lua 文件
                # 只清理SIMION生成的临时文件（不显示详细日志）
                self._cleanup_simion_temp_files(verbose=False)
                print(f"工作空间临时文件已清理，目录结构已保留: {self.workspace_dir}")
            else:
                print("删除整个工作空间目录")
                # 完整删除整个临时工作目录
                shutil.rmtree(self.workspace_dir, ignore_errors=True)
                print("工作空间已完全清理")
        except Exception as e:
            print(f"清理工作空间失败: {e}")
    
    def _cleanup_simion_temp_files(self, verbose=False):
        """清理SIMION生成的临时文件"""
        try:
            # 基础临时文件扩展名（不包含PA文件）
            base_temp_extensions = [
                '.tmp', '.wrk', '.log', '.out', '.err', '.cache', '.lock'
            ]
            
            # 动态获取PA临时文件扩展名
            pa_temp_extensions = self._get_pa_temp_extensions()
            
            # 合并所有临时文件扩展名
            all_temp_extensions = base_temp_extensions + pa_temp_extensions
            
            # 清理每个run目录中的临时文件
            if self.runs_dir.exists():
                for run_dir in self.runs_dir.iterdir():
                    if run_dir.is_dir():
                        self._cleanup_directory_temp_files(run_dir, all_temp_extensions, verbose)
            
            # 清理common目录中的临时文件
            if self.common_dir.exists():
                self._cleanup_directory_temp_files(self.common_dir, all_temp_extensions, verbose)
                
        except Exception as e:
            print(f"清理SIMION临时文件失败: {e}")
    
    def _get_pa_temp_extensions(self):
        """动态获取PA临时文件扩展名"""
        pa_extensions = []
        
        try:
            # 检查原始目录中的PA文件，确定电极数量
            base_dir = Path(self.base_iob_path).parent
            iob_name = self.base_iob_path.stem
            
            # 查找所有PA文件
            pa_files = list(base_dir.glob(f"{iob_name}.pa*"))
            
            for pa_file in pa_files:
                suffix = pa_file.suffix
                if suffix and suffix not in ['.pa']:  # 排除主PA文件
                    pa_extensions.append(suffix)
            
            # 如果没有找到PA文件，使用默认的
            if not pa_extensions:
                pa_extensions = ['.pa#', '.pa0', '.pa1', '.pa2']
            
            return pa_extensions
            
        except Exception as e:
            print(f"获取PA临时文件扩展名失败: {e}")
            # 返回默认的PA扩展名
            return ['.pa#', '.pa0', '.pa1', '.pa2']
    
    def get_electrode_count(self):
        """获取电极数量"""
        try:
            base_dir = Path(self.base_iob_path).parent
            iob_name = self.base_iob_path.stem
            
            # 查找所有PA文件
            pa_files = list(base_dir.glob(f"{iob_name}.pa*"))
            
            # 计算电极数量（排除主PA文件）
            electrode_count = len([f for f in pa_files if f.suffix != '.pa'])
            
            print(f"检测到电极数量: {electrode_count}")
            return electrode_count
            
        except Exception as e:
            print(f"获取电极数量失败: {e}")
            return 0
    
    def _cleanup_directory_temp_files(self, directory, temp_extensions, verbose=False):
        """清理指定目录中的临时文件"""
        try:
            for file_path in directory.iterdir():
                if file_path.is_file():
                    # 检查是否是临时文件
                    if any(file_path.suffix.lower() == ext for ext in temp_extensions):
                        # 对于PA文件，只删除临时文件，保留原始文件
                        if file_path.suffix.lower().startswith('.pa'):
                            if self._is_pa_temp_file(file_path):
                                try:
                                    file_path.unlink()
                                    if verbose:
                                        print(f"删除PA临时文件: {file_path.name}")
                                except Exception as e:
                                    if verbose:
                                        print(f"删除PA临时文件失败 {file_path.name}: {e}")
                        else:
                            # 其他临时文件直接删除
                            try:
                                file_path.unlink()
                                if verbose:
                                    print(f"删除临时文件: {file_path.name}")
                            except Exception as e:
                                if verbose:
                                    print(f"删除临时文件失败 {file_path.name}: {e}")
                    # 检查是否是SIMION生成的临时文件（以特定前缀命名）
                    elif file_path.name.startswith(('temp_', 'simion_', 'fly_')):
                        try:
                            file_path.unlink()
                            if verbose:
                                print(f"删除SIMION临时文件: {file_path.name}")
                        except Exception as e:
                            if verbose:
                                print(f"删除SIMION临时文件失败 {file_path.name}: {e}")
        except Exception as e:
            if verbose:
                print(f"清理目录临时文件失败 {directory}: {e}")
    
    def _is_pa_temp_file(self, file_path):
        """判断是否是PA临时文件（而非原始PA文件）"""
        try:
            # 检查文件是否在common目录中（原始文件）
            if file_path.parent == self.common_dir:
                return False  # common目录中的PA文件是原始文件，不删除
            
            # 检查文件是否是链接文件
            if file_path.is_symlink() or (platform.system() == "Windows" and self._is_hardlink(file_path)):
                return False  # 链接文件指向原始文件，不删除
            
            # 其他情况认为是临时文件
            return True
            
        except Exception:
            return True  # 出错时保守处理，认为是临时文件
    
    def _is_hardlink(self, file_path):
        """检查文件是否是硬链接（简化版本）"""
        try:
            # 简化版本：检查文件是否在run目录中且不是符号链接
            # 如果文件在run目录中，且不是符号链接，则认为是硬链接（指向原始文件）
            if file_path.parent.parent == self.runs_dir:  # 在runX目录中
                return True  # 认为是硬链接，不删除
            return False
        except Exception:
            return False
    
    def get_workspace_info(self):
        """获取工作空间信息"""
        return {
            'workspace_dir': str(self.workspace_dir),
            'common_dir': str(self.common_dir),
            'runs_dir': str(self.runs_dir),
            'thread_count': self.thread_count,
            'base_iob': str(self.base_iob_path),
            'electrode_count': self.get_electrode_count(),
            'pa_temp_extensions': self._get_pa_temp_extensions()
        }


def create_simion_workspace(iob_path, thread_count=None):
    """
    创建SIMION工作空间的便捷函数
    
    参数:
        iob_path: IOB文件路径
        thread_count: 线程数量
        
    返回:
        SimionWorkspaceManager实例
    """
    manager = SimionWorkspaceManager(iob_path, thread_count)
    if manager.setup_workspace():
        return manager
    else:
        return None
